#!/usr/bin/env python3

import os

import anthropic

from memori import Memori
from tests.database.core import TestDBSession

if os.environ.get("ANTHROPIC_API_KEY", None) is None:
    raise RuntimeError("ANTHROPIC_API_KEY is not set")

os.environ["MEMORI_TEST_MODE"] = "1"

session = TestDBSession
client = anthropic.Anthropic()

mem = Memori(conn=session).llm.register(client)

# Multiple registrations should not cause an issue.
mem.llm.register(client)

mem.attribution(entity_id="123", process_id="456")

print("-" * 25)

query = "What color is the planet Mars?"
print(f"me: {query}")

response = client.messages.create(
    model="claude-sonnet-4-20250514",
    max_tokens=1024,
    messages=[{"role": "user", "content": query}],
)

print("-" * 25)

print(f"llm: {response.content[0].text}")

print("-" * 25)

query = "That planet we're talking about, in order from the sun which one is it?"
print(f"me: {query}")

print("-" * 25)
print("CONVERSATION INJECTION OCCURRED HERE!\n")

response = client.messages.create(
    model="claude-sonnet-4-20250514",
    max_tokens=1024,
    messages=[{"role": "user", "content": query}],
)

print("-" * 25)

print(f"llm: {response.content[0].text}")

print("-" * 25)
